import sys
import random
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QLabel, QVBoxLayout
from PyQt5.QtGui import QPainter, QColor, QFont, QKeyEvent
from PyQt5.QtCore import Qt, QTimer, QRect

class SnakeGame(QWidget):
    def __init__(self):
        super().__init__()
        self.setFixedSize(600, 600)
        self.setStyleSheet("background-color: #111111;")
        self.setWindowTitle('PyQt5贪吃蛇游戏')
        
        # 游戏区域设置
        self.grid_size = 20
        self.grid_width = self.width() // self.grid_size
        self.grid_height = self.height() // self.grid_size
        
        # 初始化游戏状态
        self.reset_game()
        
        # 游戏计时器
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update_game)
        self.timer.start(150)  # 控制游戏速度
        
    def reset_game(self):
        """重置游戏状态"""
        # 蛇的初始位置和方向
        self.snake = [(self.grid_width // 2, self.grid_height // 2)]
        self.direction = Qt.Key_Right
        
        # 初始食物位置
        self.food = self.generate_food()
        
        # 游戏状态和分数
        self.game_over = False
        self.score = 0
        
    def generate_food(self):
        """生成食物位置"""
        while True:
            food = (random.randint(0, self.grid_width - 1), 
                    random.randint(0, self.grid_height - 1))
            # 确保食物不会出现在蛇身上
            if food not in self.snake:
                return food
    
    def update_game(self):
        """更新游戏状态"""
        if self.game_over:
            return
            
        # 根据当前方向计算新蛇头位置
        head_x, head_y = self.snake[0]
        if self.direction == Qt.Key_Up:
            head_y -= 1
        elif self.direction == Qt.Key_Down:
            head_y += 1
        elif self.direction == Qt.Key_Left:
            head_x -= 1
        elif self.direction == Qt.Key_Right:
            head_x += 1
            
        # 检查是否撞墙
        if (head_x < 0 or head_x >= self.grid_width or 
            head_y < 0 or head_y >= self.grid_height):
            self.game_over = True
            return
            
        # 检查是否撞到自己
        if (head_x, head_y) in self.snake:
            self.game_over = True
            return
            
        # 移动蛇
        self.snake.insert(0, (head_x, head_y))
        
        # 检查是否吃到食物
        if (head_x, head_y) == self.food:
            # 吃到食物，生成新食物
            self.food = self.generate_food()
            self.score += 10
        else:
            # 没吃到食物，移除蛇尾
            self.snake.pop()
            
        self.update()  # 触发重绘
    
    def paintEvent(self, event):
        """绘制游戏元素"""
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        
        # 绘制食物
        food_x, food_y = self.food
        food_rect = QRect(
            food_x * self.grid_size, 
            food_y * self.grid_size,
            self.grid_size, 
            self.grid_size
        )
        painter.setBrush(QColor(255, 50, 50))  # 红色
        painter.drawEllipse(food_rect)
        
        # 绘制蛇
        for i, (x, y) in enumerate(self.snake):
            snake_rect = QRect(
                x * self.grid_size + 1, 
                y * self.grid_size + 1,
                self.grid_size - 2, 
                self.grid_size - 2
            )
            # 蛇头用不同颜色
            if i == 0:
                painter.setBrush(QColor(0, 180, 0))  # 绿色蛇头
            else:
                painter.setBrush(QColor(0, 200, 0))  # 亮绿色蛇身
            painter.drawRect(snake_rect)
        
        # 绘制分数
        painter.setPen(QColor(255, 255, 255))
        painter.setFont(QFont('Arial', 12))
        painter.drawText(10, 20, f"分数: {self.score}")
        
        # 游戏结束提示
        if self.game_over:
            painter.setPen(QColor(255, 50, 50))
            painter.setFont(QFont('Arial', 30))
            painter.drawText(self.rect(), Qt.AlignCenter, "游戏结束!")
            painter.setFont(QFont('Arial', 16))
            painter.drawText(
                self.rect().adjusted(0, 40, 0, 0), 
                Qt.AlignCenter, 
                f"最终得分: {self.score}"
            )
            painter.setFont(QFont('Arial', 12))
            painter.drawText(
                self.rect().adjusted(0, 70, 0, 0), 
                Qt.AlignCenter, 
                "按R键重新开始"
            )
    
    def keyPressEvent(self, event: QKeyEvent):
        """处理键盘事件"""
        key = event.key()
        
        # 游戏结束时按R键重新开始
        if self.game_over and key == Qt.Key_R:
            self.reset_game()
            return
            
        # 忽略游戏结束时的按键
        if self.game_over:
            return
            
        # 防止180度转向
        if (key == Qt.Key_Up and self.direction != Qt.Key_Down or
            key == Qt.Key_Down and self.direction != Qt.Key_Up or
            key == Qt.Key_Left and self.direction != Qt.Key_Right or
            key == Qt.Key_Right and self.direction != Qt.Key_Left):
            self.direction = key

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('PyQt5贪吃蛇游戏')
        self.setFixedSize(600, 600)
        
        # 创建游戏主界面
        self.game = SnakeGame()
        self.setCentralWidget(self.game)
        
        # 游戏说明
        self.statusBar().showMessage("使用方向键控制蛇的移动，按R键重新开始游戏")
        
    def keyPressEvent(self, event):
        # 将键盘事件传递给游戏组件
        self.game.keyPressEvent(event)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())